import sensor, image, utime, lcd
import KPU as kpu
import gc, sys, os
from machine import UART,Timer
from fpioa_manager import fm

mask_count = 1
no_mask_count = 1

color_R = (255, 0, 0)
color_G = (0, 255, 0)
color_B = (0, 0, 255)

def drawConfidenceText(image, rol, classid, value):
    text = ""
    _confidence = int(value * 100)

    if classid == 1:
        text = 'mask: ' + str(_confidence) + '%'
        color_text=color_G
    else:
        text = 'no_mask: ' + str(_confidence) + '%'
        color_text=color_R
    image.draw_string(rol[0], rol[1], text, color=color_text, scale=2.5)

class uart():
    def __init__(self, name):
        self.type = None
        self.data = None
        self.name = name
        self.name = UART(UART.UART1, 2400, read_buf_len=4096)
    def uart_init(self):
        # 映射串口引脚
        fm.register(32, fm.fpioa.UART1_RX, force=True)
        fm.register(33, fm.fpioa.UART1_TX, force=True)
        return self.name

    def uart_seed(self,input,type):
        self.type = type
        self.data = input
        if self.type == True:
            self.name.write('EE{}\n'.format(self.data))
        else:
            self.name.write('AA{}\n'.format(self.data))

def save_image(img0):
    os.chdir("/sd")
    images_dir = "cap_images"
    if images_dir not in os.listdir():
        os.mkdir(images_dir)
    global no_mask_count
    f_name = "{}/no_mask-{}.jpg".format(images_dir, no_mask_count)
    img0.save(f_name, quality=95)
    Uar.uart_seed(input=no_mask_count, type=False)
    no_mask_count += 1


def main(anchor, model_addr):
    lcd.init()
    sensor.reset()
    sensor.set_pixformat(sensor.RGB565)
    sensor.set_framesize(sensor.QVGA)
    sensor.set_hmirror(0)
    sensor.run(1)
    task = kpu.load(model_addr)
    flag = 0
    _ = kpu.init_yolo2(task, 0.5, 0.3, 5, anchor)
    try:
        while True:
            img = sensor.snapshot()
            code = kpu.run_yolo2(task, img)
            if code:
                totalRes = len(code)

                for item in code:
                    confidence = float(item.value())
                    itemROL = item.rect()
                    classID = int(item.classid())

                    if classID == 0 and confidence > 0.5 and flag != 1:
                       if totalRes == 1:
                          print("图片保存{}".format(no_mask_count))
                          save_image(img)
                          flag = 1
                          _ = img.draw_rectangle(itemROL, color=color_R, tickness=5)
                          drawConfidenceText(img, (0, 0), 1, confidence)
                          continue

                    if classID == 1 and confidence > 0.75 and flag != 2:
                        flag = 2
                        global mask_count
                        Uar.uart_seed(input=mask_count, type=True)
                        print("mask_count:{}\r\n".format(mask_count))
                        mask_count += 1
                        _ = img.draw_rectangle(itemROL, color_G, tickness=5)
                        drawConfidenceText(img, (0, 0), 0, confidence)
                        continue

                    if confidence < 0.52 and classID == 0:
                        flag = 3

                        continue
            _ = lcd.display(img)
    except Exception as e:
        raise e
    finally:
        kpu.deinit(task)
        main(anchor=anchor, model_addr="/sd/mask.smodel")


if __name__ == "__main__":
    try:
        Uar = uart(uart)
        Uar.uart_init()
        labels = ['no_mask', 'mask']
        anchors = (0.1606, 0.3562, 0.4712, 0.9568, 0.9877, 1.9108, 1.8761, 3.5310, 3.4423, 5.6823)
        main(anchor=anchors, model_addr="/sd/mask.smodel")
    except Exception as e:
        raise e
    finally:
        gc.collect()
